""" This script reads a file of a fullerene structure in POSCAR or CONTCAR 
format, and ouputs the fullerene's radius in Angstrom.
It can read compressed files of .gz (or .gzip) format.
This script is to process double-shell fullerene molecules, with inner shell's
radius less than 5 angstrom. Example molecules are C60atC240, C80atC180.
"""
import sys
import os.path
import gzip    # gzip is a core module, so this should be portable.
import math

def usage():
    """ Show help message."""
    print("Usage: python3 fullerene-radius.py [CONTCAR]")
    print("Note: If no file is specified, default is CONTCAR or CONTCAR.gz.")
    print("You have to specify POSCAR if it is the desired input file.")

# Parse the command line arguments.
file_structure = ""
if len(sys.argv)==1:
    for f in ["CONTCAR", "CONTCAR.gz", "CONTCAR.gzip"]:
        if os.path.isfile(f): 
            file_structure = f
            break
    if file_structure=="":
        # No file is specified, and CONTCAR file is missing. This is an error.
        print("I can't find CONTCAR file. You need to specify the filename.")
        usage()
        sys.exit(1)
elif len(sys.argv)==2:
    file_structure = sys.argv[1]
else:
    usage()
    sys.exit(1)

# Open the file for reading.
extension_str = os.path.splitext(file_structure)[1]
if extension_str.upper() in [".GZ", ".GZIP"]:
    f = gzip.open(file_structure, "rb")
    all_lines = [x.decode('utf-8') for x in f.readlines()]  
    # I need this decode to suppress IO errors
else:
    f = open(file_structure, "r") 
    all_lines = f.readlines()
f.close()

# Assume the basis is [a, 0, 0; 0, b, 0; 0, 0, c].
# TO DO: warn if the basis is not so.
scale_factor = float(all_lines[1])
latt_a = float(all_lines[2].split()[0]) * scale_factor
latt_b = float(all_lines[3].split()[1]) * scale_factor
latt_c = float(all_lines[4].split()[2]) * scale_factor

lines_to_skip = 0
# skip a correct number of lines
if all_lines[5].strip()[0].isalpha():    # element's name
    lines_to_skip = lines_to_skip + 1
n_atoms = [int(x) for x in all_lines[5+lines_to_skip].split()]
if all_lines[6+lines_to_skip][0] in ['S', 's']: # selective dynamics
    lines_to_skip = lines_to_skip + 1

letter0 = all_lines[6+lines_to_skip][0]
if letter0 in ['c', 'C']: # cartesian
    #print("The cartesian coordinate system is being used.")
    coordinate_system = 0
elif letter0 in ['d', 'D']: # direct
    #print("The direct coordinate system is being used.")
    coordinate_system = 1
else:
    print("***I can't understand the coordinate system. Corrupted file?")
    sys.exit(2)

# find the mass center
coord = []
mc_x = 0.0
mc_y = 0.0
mc_z = 0.0
if coordinate_system==1:
    for line in all_lines[7+lines_to_skip : ]:
        if len(line)<5: break  # an empty line indicates end of coordinates.
        xyz = [float(x) for x in line.split()[0:3]]
        tx = xyz[0]*latt_a
        ty = xyz[1]*latt_b
        tz = xyz[2]*latt_c
        mc_x = mc_x + tx
        mc_y = mc_y + ty
        mc_z = mc_z + tz
        coord.append([tx, ty, tz])
else:
    for line in all_lines[7+lines_to_skip : ]:
        if len(line)<5: break  # an empty line indicates end of coordinates.
        xyz = [float(x) for x in line.split()[0:3]]
        mc_x = mc_x + xyz[0]
        mc_y = mc_y + xyz[1]
        mc_z = mc_z + xyz[2]
        coord.append(xyz)

n = len(coord)
mc_x = mc_x/n
mc_y = mc_y/n
mc_z = mc_z/n

# calculates average distance between atoms and mass-center
r_inner = 0.0
r_outer = 0.0
n_inner = 0
n_outer = 0
for xyz in coord:
    tx = xyz[0]-mc_x
    ty = xyz[1]-mc_y
    tz = xyz[2]-mc_z
    r = math.sqrt(tx*tx + ty*ty + tz*tz)
    if r<5.0:
        r_inner = r_inner + r
        n_inner = n_inner + 1
    else:
        r_outer = r_outer + r
        n_outer = n_outer + 1

# Show the result
if n_inner==0:
    print("Only one shell is found. It should be a single-shell molecule.")
    r_outer = r_outer/n_outer
    print("Number of atoms: {0:d}. Radius: {1:f}".format(n_outer, r_outer))
elif n_outer==0:
    print("Only one shell is found. It should be a single-shell molecule.")
    r_inner = r_inner/n_inner
    print("Number of atoms: {0:d}. Radius: {1:f}".format(n_inner, r_inner))
else:
    r_inner = r_inner/n_inner
    r_outer = r_outer/n_outer 
    print("Amond {0:d} atoms, {1:d} compose the inner shell,"
          "and {2:d} compose the outer one.".format(n, n_inner, n_outer))
    print("Radii of inner and outer shells:", r_inner, r_outer)

